package com.ruoyi.framework.security.handle;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONObject;
import com.ruoyi.common.constant.Constants;
import com.ruoyi.common.core.domain.AjaxResult;
import com.ruoyi.common.core.domain.model.LoginUser;
import com.ruoyi.common.utils.MessageUtils;
import com.ruoyi.common.utils.ServletUtils;
import com.ruoyi.framework.config.OAuth2SecurityConfig;
import com.ruoyi.framework.config.properties.KeycloakConfig;
import com.ruoyi.framework.manager.AsyncManager;
import com.ruoyi.framework.manager.factory.AsyncFactory;
import com.ruoyi.framework.web.service.TokenService;
import jakarta.servlet.ServletException;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import org.apache.commons.lang3.StringUtils;
import org.springframework.security.core.Authentication;
import org.springframework.security.oauth2.client.authentication.OAuth2AuthenticationToken;
import org.springframework.security.oauth2.client.registration.ClientRegistration;
import org.springframework.security.oauth2.client.registration.ClientRegistrationRepository;
import org.springframework.security.oauth2.core.oidc.user.OidcUser;
import org.springframework.security.oauth2.server.resource.authentication.JwtAuthenticationToken;
import org.springframework.security.web.authentication.logout.SimpleUrlLogoutSuccessHandler;
import org.springframework.security.web.util.UrlUtils;
import org.springframework.util.Assert;
import org.springframework.web.util.UriComponents;
import org.springframework.web.util.UriComponentsBuilder;

import java.io.IOException;
import java.net.URI;
import java.nio.charset.StandardCharsets;
import java.util.Collections;

/**
 * 自定义登出成功处理器，用于处理OIDC身份认证的单点登出流程。
 * 重写了determineTargetUrl方法，根据身份认证类型获取对应的登出端点URI。
 */
public class OAuth2LogoutSuccessHandler extends SimpleUrlLogoutSuccessHandler {

    private final ClientRegistrationRepository clientRegistrationRepository;
    private final KeycloakConfig keycloakConfig;
    private String postLogoutRedirectUri;
    private final TokenService tokenService;

    public OAuth2LogoutSuccessHandler(ClientRegistrationRepository clientRegistrationRepository,
                                      KeycloakConfig keycloakConfig, TokenService tokenService) {
        Assert.notNull(clientRegistrationRepository, "ClientRegistrationRepository must not be null");
        this.clientRegistrationRepository = clientRegistrationRepository;
        this.keycloakConfig = keycloakConfig;
        this.tokenService = tokenService;
    }

    // 因为同时使用了身份认证和资源管理，因此需要判断authentication的类型，根据类型获取对应的登出端点和idToken
    protected String determineTargetUrl(HttpServletRequest request, HttpServletResponse response, Authentication authentication) {
        String targetUrl = null;
        if (authentication instanceof OAuth2AuthenticationToken && authentication.getPrincipal() instanceof OidcUser) {
            String registrationId = ((OAuth2AuthenticationToken)authentication).getAuthorizedClientRegistrationId();
            ClientRegistration clientRegistration = this.clientRegistrationRepository.findByRegistrationId(registrationId);
            URI endSessionEndpoint = this.endSessionEndpoint(clientRegistration);
            if (endSessionEndpoint != null) {
                String idToken = this.idToken(authentication);
                URI postLogoutRedirectUri = this.postLogoutRedirectUri(request);
                targetUrl = this.endpointUri(endSessionEndpoint, idToken, postLogoutRedirectUri);
            }
        } else if(authentication instanceof JwtAuthenticationToken){
            ClientRegistration clientRegistration = this.clientRegistrationRepository.findByRegistrationId(keycloakConfig.getProvider());
            URI endSessionEndpoint = this.endSessionEndpoint(clientRegistration);
            if (endSessionEndpoint != null) {
                String idToken = (String) request.getAttribute("idToken");
                URI postLogoutRedirectUri = this.postLogoutRedirectUri(request);
                targetUrl = this.endpointUri(endSessionEndpoint, idToken, postLogoutRedirectUri);
            }
        }
        if (targetUrl == null) {
            targetUrl = super.determineTargetUrl(request, response);
        }

        return targetUrl;
    }

    private URI endSessionEndpoint(ClientRegistration clientRegistration) {
        URI result = null;
        if (clientRegistration != null) {
            Object endSessionEndpoint = clientRegistration.getProviderDetails().getConfigurationMetadata().get("end_session_endpoint");
            if (endSessionEndpoint != null) {
                result = URI.create(endSessionEndpoint.toString());
            }
        }

        return result;
    }

    private String idToken(Authentication authentication) {
        return ((OidcUser)authentication.getPrincipal()).getIdToken().getTokenValue();
    }

    private URI postLogoutRedirectUri(HttpServletRequest request) {
        if (this.postLogoutRedirectUri == null) {
            return null;
        } else {
            UriComponents uriComponents = UriComponentsBuilder.fromHttpUrl(UrlUtils.buildFullRequestUrl(request)).replacePath(request.getContextPath()).replaceQuery(null).fragment(null).build();
            return UriComponentsBuilder.fromUriString(this.postLogoutRedirectUri).buildAndExpand(Collections.singletonMap("baseUrl", uriComponents.toUriString())).toUri();
        }
    }

    private String endpointUri(URI endSessionEndpoint, String idToken, URI postLogoutRedirectUri) {
        UriComponentsBuilder builder = UriComponentsBuilder.fromUri(endSessionEndpoint);
        builder.queryParam("id_token_hint", idToken);
        if (postLogoutRedirectUri != null) {
            builder.queryParam("post_logout_redirect_uri", postLogoutRedirectUri);
        }

        return builder.encode(StandardCharsets.UTF_8).build().toUriString();
    }

    public void setPostLogoutRedirectUri(String postLogoutRedirectUri) {
        Assert.notNull(postLogoutRedirectUri, "postLogoutRedirectUri cannot be null");
        this.postLogoutRedirectUri = postLogoutRedirectUri;
    }

    @Override
    public void onLogoutSuccess(HttpServletRequest request, HttpServletResponse response, Authentication authentication) throws IOException, ServletException {
        // 让系统原本的登录 token 失效
        try{
            LoginUser loginUser = tokenService.getLoginUser(request);
            if (com.ruoyi.common.utils.StringUtils.isNotNull(loginUser)) {
                String userName = loginUser.getUsername();
                String token = loginUser.getToken();
                tokenService.delLoginUser(token);
                AsyncManager.me().execute(AsyncFactory.recordLogininfor(userName, Constants.LOGOUT, MessageUtils.message("user.logout.success")));
                OAuth2SecurityConfig.tokenMap.values().removeIf(value -> value.compareTo(token) == 0);
            }
        }catch (Exception e){
            logger.error(e.getMessage(), e);
        }

        // 让 oauth2 登录的会话失效
        try{
            String sessionId = (String) request.getAttribute("sessionId");
            if(StringUtils.isNotBlank(sessionId)){
                OAuth2SecurityConfig.sessionMap.values().removeIf(value -> value.compareTo(sessionId) == 0);
            }
        }catch (Exception e){
            logger.error(e.getMessage(), e);
        }

        // 将要注销成功后要重定向的地址返回给前端，由前端去跳转，避免跨域问题
        String targetUrl = this.determineTargetUrl(request, response, authentication);
        JSONObject responseData = new JSONObject();
        responseData.put("endpoint", targetUrl);
        ServletUtils.renderString(response, JSON.toJSONString(AjaxResult.success(MessageUtils.message("user.logout.success"), responseData)));
    }
}
